import torch
import torch.nn as nn
import torch.nn.functional as F
def BCELoss(pred, mask):
    return F.binary_cross_entropy_with_logits(pred, mask)
def IOULoss(pred, mask):
    pred  = torch.sigmoid(pred)
    inter = (pred*mask).sum(dim=(2,3))
    union = (pred+mask).sum(dim=(2,3))
    iou  = 1-(inter+1)/(union-inter+1)
    return iou.mean()
class EdgeSaliencyLoss(nn.Module):
    def __init__(self, device):
        super(EdgeSaliencyLoss, self).__init__()

        # self.alpha_sal = alpha_sal

        self.laplacian_kernel = torch.tensor([[-1., -1., -1.], [-1., 8., -1.], [-1., -1., -1.]], dtype=torch.float, requires_grad=False)
        self.laplacian_kernel = self.laplacian_kernel.view((1, 1, 3, 3))  # Shape format of weight for convolution
        self.laplacian_kernel = self.laplacian_kernel.to(device)

    def forward(self, y_pred, y_gt):
        # Generate edge maps
        y_gt_edges = F.relu(torch.tanh(F.conv2d(y_gt, self.laplacian_kernel, padding=(1, 1))))
        y_pred_edges = F.relu(torch.tanh(F.conv2d(y_pred, self.laplacian_kernel, padding=(1, 1))))
        edge_loss = F.binary_cross_entropy(input=y_pred_edges, target=y_gt_edges)
        return edge_loss